import cv2
import matplotlib.pyplot as plt
import os 
import sys
import numpy as np
import json
os.chdir(sys.path[0])

def json_to_label(json_path, label_save_path,  path_allclass_txt, label_size=(224, 224)):
    """
    该函数直接由.json文件生成单通道8位标签图像,各类位置区域像素对应各类序号0,1,2,3,4......等
    :param json_path: 包含.json文件的文件夹路径
    :param label_save_path: 标签存放位置
    :param path_allclass_txt: 包含所有类别的txt文件路径
    :param label_size: 标签图像的大小
    :return:
    """
    class_txt = open(path_allclass_txt, "r")
    category_types = class_txt.read().splitlines()
    print('All class name:\n', category_types)

    for file in os.listdir(json_path):
        mask = np.zeros([label_size[1], label_size[0], 1], np.uint8)  # 创建一个大小和原图相同的空白图像
        with open(os.path.join(json_path, file), "r") as f:
            label = json.load(f)
        shapes = label["shapes"]
        for shape in shapes:
            category = shape["label"]
            points = shape["points"]
            points_array = np.array(points, dtype=np.int32)# 填充
            mask = cv2.fillPoly(mask, [points_array], 1+category_types.index(category)) #在对应位置填充类别序号
        name = file.split('.')[0]
        cv2.imwrite(label_save_path + '/' + name + ".png", mask)
    print('json file process ok, get %d label images' % len(os.listdir(json_path)))

def normalize_label(x,y,w,h):
    """像素坐标归一化

    """
    if isinstance(x, str):
        x = float(x)
        y = float(y)
    x1=x/w
    y1=y/h
    return str(x1),str(y1)


def json2txt(json_path,txt_path):
    """由labelme生成的json标签  生成   txt标签  (只保留类别0)

    Args:
        json_path (str): json地址
        txt_path (str): txt保存地址
    
    """
    class_txt = open(path_allclass_txt, "r")
    category_types = class_txt.read().splitlines()
    print('All class name:\n', category_types)
    
    for file in os.listdir(json_path):
        name = file.split('.')[0]
        txt_name=txt_path + name.zfill(12) + ".txt"
        
        T=open(txt_name,"a+")
        with open(os.path.join(json_path, file), "r") as f:
            label = json.load(f)
        shapes = label["shapes"]
        W=label["imageWidth"]
        H=label["imageHeight"]
        for shape in shapes:
            category = shape["label"]
            idx=category_types.index(category)
            if idx!=0:
                # 只要标签不是第一个,就跳过.
                continue
            T.write(str(idx)+ " ")
            points = shape["points"] #list
            for point in points:
                x,y=normalize_label(point[0],point[1],W,H)
                x=cut_str_by_num(x,6) #取小数点后6位并转换为字符串
                y=cut_str_by_num(y,6)
                T.write(x+ " ")
                T.write(y+ " ")
            T.write("\n")
        T.close()
        
def cut_str_by_num(str_,num_):
    """截取str的小数点后num_位

    Args:
        str_ (str): float数值的str 例如 "0.11133"
        num_ (int): 保存的小数点位数
    Return:
        str_
    """
    str_list=str_.split(".")
    s1_new = str_list[0] + '.' + str_list[1][:num_]
    return s1_new

def bit24_to_bit8(bit24_path, bit8_path):
    """从黑白图转换到标签图

    """
    for file in os.listdir(bit24_path):
        img1 = cv2.imread(bit24_path + '/' + file, -1)
        print(img1.shape)
        img = img1
        img[img != 0] = 255  # label use
        plt.subplot(1, 3, 1)
        plt.imshow(img1)
        plt.title('img')
        plt.subplot(1, 3, 2)
        plt.imshow(img, 'gray')
        plt.show()
        img_name = file.split('.')[0]
        # cv2.imwrite(bit8_path + '/' + img_name + '.png', img)

def img_change_name(image_path):
    """把图片名字变成12位数 例如 000000000000.jpg

    """
    for file in os.listdir(image_path):
        name = file.split('.')[0]
        file_new_name=image_path + name.zfill(12) + ".jpg"
        os.rename(image_path+file,file_new_name)

if __name__=="__main__":
    json_path = '../../datasets/xwyd_10x10_road/json/'
    # label_save_path= '../../datasets/xwyd_10x10_road/json_img/'
    # path_allclass_txt = '../../datasets/xwyd_10x10_road/class_name.txt'
    # label_show_path='../../datasets/xwyd_10x10_road/json_img_white/'
    yoloseg_txt_label='../../datasets/xwyd_10x10_road/labels/origin/'
    # image_path="../../datasets/xwyd_10x10_road/images/"
    # json_to_label(json_path,label_save_path,path_allclass_txt) # 生成语义标签
    # bit24_to_bit8(label_save_path,label_show_path) #可视化语义标签
    json2txt(json_path,yoloseg_txt_label)
    # img_change_name(image_path) 

    ## <<<<test:cut_str_by_num
    # str1="1.12345678"
    # str2=cut_str_by_num(str1,5)
    # print(str2)
    ## >>>>>test:cut_str_by_num
